package cc.shacocloud.mirage.maven.plugin;

import cc.shacocloud.mirage.utils.CacheURLClassLoader;
import cc.shacocloud.mirage.utils.annotation.AnnotatedElementUtils;
import cc.shacocloud.mirage.utils.reflection.ReflectUtil;
import org.apache.commons.compress.archivers.jar.JarArchiveEntry;
import org.apache.maven.artifact.Artifact;
import org.apache.maven.execution.MavenSession;
import org.apache.maven.plugin.MojoExecutionException;
import org.apache.maven.plugin.MojoFailureException;
import org.apache.maven.plugins.annotations.LifecyclePhase;
import org.apache.maven.plugins.annotations.Mojo;
import org.apache.maven.plugins.annotations.Parameter;
import org.apache.maven.plugins.annotations.ResolutionScope;
import org.apache.maven.project.MavenProject;
import org.jetbrains.annotations.NotNull;
import org.jetbrains.annotations.Nullable;

import java.io.*;
import java.lang.annotation.Annotation;
import java.lang.reflect.Method;
import java.nio.file.attribute.FileTime;
import java.time.OffsetDateTime;
import java.util.ArrayList;
import java.util.Enumeration;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.TimeUnit;
import java.util.jar.Attributes;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.jar.Manifest;
import java.util.stream.Collectors;

/**
 * 重新打包现有的 JAR，以便可以使用 {@code java -jar} 启动 java 程序
 *
 * @author 思追(shaco)
 */
@Mojo(name = "repackage", defaultPhase = LifecyclePhase.PACKAGE, threadSafe = true,
        requiresDependencyResolution = ResolutionScope.COMPILE_PLUS_RUNTIME, requiresDependencyCollection = ResolutionScope.COMPILE_PLUS_RUNTIME)
public class RepackageMojo extends AbstractLoggingMojo {
    
    public static final String MAIN_CLASS_ATTRIBUTE = "Main-Class";
    public static final String MAIN_CLASS = "cc.shacocloud.mirage.loader.MirageStarter";
    
    public static final String START_CLASS_ATTRIBUTE = "Mirage-Start-Class";
    
    public static final String VERSION_ATTRIBUTE = "Mirage-Version";
    
    public static final String CLASSES_ATTRIBUTE = "Mirage-Classes";
    public static final String CLASSES_LOCATION = "BOOT-INF/classes/";
    
    public static final String LIB_ATTRIBUTE = "Mirage-Lib";
    public static final String LIB_LOCATION = "BOOT-INF/lib/";
    
    public static final String LIB_INDEX_ATTRIBUTE = "Mirage-Lib-Index";
    public static final String LIB_INDEX_LOCATION = "BOOT-INF/artifactLib.index";
    
    public static final String MAIN_CLASS_ANNOTATION = "cc.shacocloud.mirage.starter.MirageBootApplication";
    public static final String LOADER_LIB_NAME = "mirage-loader-" + RepackageMojo.class.getPackage().getImplementationVersion() + ".jar";
    
    /**
     * 启动类型配置，全类名
     * <p>
     * 如果未设置，将进行扫描
     */
    @Parameter(name = "mainClass")
    protected String mainClass;
    
    /**
     * 输出的文件路径
     */
    @Parameter(defaultValue = "${project.build.directory}", required = true)
    protected File outputDirectory;
    
    /**
     * 生成的 jar 名称
     */
    @Parameter(defaultValue = "${project.build.finalName}", readonly = true)
    protected String finalName;
    @Parameter(defaultValue = "${project}", required = true, readonly = true)
    protected MavenProject project;
    @Parameter(defaultValue = "${session}", readonly = true, required = true)
    protected MavenSession session;
    /**
     * 输出的时间戳
     */
    @Parameter(defaultValue = "${project.build.outputTimestamp}")
    private String outputTimestamp;
    
    @Override
    public void execute() throws MojoExecutionException, MojoFailureException {
        if (project.getPackaging().equals("pom")) {
            logWarn(() -> "无法打包类型为 pom 的项目！");
            return;
        }
        
        try {
            doPackage();
        } catch (Exception e) {
            throw new MojoExecutionException(e);
        }
    }
    
    /**
     * 执行打包
     */
    protected void doPackage() throws Exception {
        File source = this.project.getArtifact().getFile().getAbsoluteFile();
        File target = getTargetFile(finalName, outputDirectory);
        
        // 如果源jar文件和输出的目标文件一致，则将其重命名为 源名称.original 的文件中
        File original = source;
        if (source.equals(target)) {
            original = new File(source.getParentFile(), source.getName() + ".original");
            original.delete();
            
            if (!source.renameTo(original)) {
                throw new IllegalStateException("无法重命名文件：" + source + " -> " + original);
            }
        }
        target.delete();
        
        // 构建 jar
        try (JarFile sourceJar = new JarFile(original)) {
            try (JarWriter jarWriter = new JarWriter(target, parseOutputTimestamp())) {
                
                // 写入清单
                jarWriter.writeManifest(buildManifest(sourceJar));
                
                // 复制源jar的内容
                jarWriter.copyByJarFile(sourceJar, new RepackageSourceJarEntryTransformer());
                
                // 依赖项转换器
                RepackageArtifactLibEntryTransformer artifactLibEntryTransformer = new RepackageArtifactLibEntryTransformer();
                
                // 依赖项索引文件
                List<ArtifactLib> artifactLibs = getLibs();
                try (ByteArrayOutputStream libIndexStream = new ByteArrayOutputStream();
                     PrintWriter libIndexWriter = new PrintWriter(new OutputStreamWriter(libIndexStream))) {
                    
                    // 添加依赖项
                    for (ArtifactLib artifactLib : artifactLibs) {
                        
                        JarArchiveEntry jarArchiveEntry = jarWriter.writeArtifactLib(artifactLib, artifactLibEntryTransformer);
                        
                        if (Objects.nonNull(jarArchiveEntry)) {
                            libIndexWriter.println(jarArchiveEntry.getName());
                        }
                    }
                    
                    // 写入索引文件
                    libIndexWriter.flush();
                    JarArchiveEntry entry = new JarArchiveEntry(LIB_INDEX_LOCATION);
                    try (ByteArrayInputStream inputStream = new ByteArrayInputStream(libIndexStream.toByteArray())) {
                        jarWriter.writeEntry(entry, new InputStreamJarEntryWriter(inputStream));
                    }
                }
                
                // 写入装载器
                ArtifactLib loaderArtifactLib = artifactLibs.stream().filter(artifactLib -> artifactLib.getName().equals(LOADER_LIB_NAME))
                        .findFirst()
                        .orElseThrow(() -> new IllegalStateException(String.format("无法获取启动装载器[%s]依赖！", LOADER_LIB_NAME)));
                try (JarFile loaderJarFile = new JarFile(loaderArtifactLib.getFile())) {
                    jarWriter.copyByJarFile(loaderJarFile, new RepackageLoaderJarEntryTransformer());
                    
                }
            }
        }
    }
    
    /**
     * 获取依赖项
     */
    @NotNull
    protected List<ArtifactLib> getLibs() {
        return project.getArtifacts()
                .stream()
                .filter(artifact -> {
                    String scope = artifact.getScope();
                    return Artifact.SCOPE_COMPILE.equals(scope) || Artifact.SCOPE_RUNTIME.equals(scope)
                            || Artifact.SCOPE_SYSTEM.equals(scope);
                })
                .filter(artifact -> Objects.nonNull(artifact.getFile()))
                .map(ArtifactLib::new)
                .collect(Collectors.toList());
    }
    
    /**
     * 项目清单
     */
    @NotNull
    protected Manifest buildManifest(@NotNull JarFile source) throws Exception {
        
        Manifest manifest;
        if (source.getManifest() != null) {
            manifest = new Manifest(source.getManifest());
        } else {
            manifest = new Manifest();
            manifest.getMainAttributes().putValue("Manifest-Version", "1.0");
        }
        
        Attributes attributes = manifest.getMainAttributes();
        
        // 在用户定义的 main class 外进行包装，由内置的启动器执行用户定义的 main class
        String mainClass = getMainClass(source, manifest);
        attributes.putValue(MAIN_CLASS_ATTRIBUTE, MAIN_CLASS);
        attributes.putValue(START_CLASS_ATTRIBUTE, mainClass);
        
        // 其他配置信息
        attributes.putValue(VERSION_ATTRIBUTE, getClass().getPackage().getImplementationVersion());
        attributes.putValue(CLASSES_ATTRIBUTE, CLASSES_LOCATION);
        attributes.putValue(LIB_ATTRIBUTE, LIB_LOCATION);
        attributes.putValue(LIB_INDEX_ATTRIBUTE, LIB_INDEX_LOCATION);
        return manifest;
    }
    
    /**
     * 获取 main 方法
     */
    @NotNull
    private String getMainClass(JarFile source, Manifest manifest) throws Exception {
        // 优先使用自定义
        if (Objects.nonNull(this.mainClass)) {
            return this.mainClass;
        }
        
        // 其次使用 jar manifest 文件中定义的
        String attributeValue = manifest.getMainAttributes().getValue(MAIN_CLASS_ATTRIBUTE);
        if (attributeValue != null) {
            return attributeValue;
        }
        
        // 最后进行扫描
        return findMainMethod(source);
    }
    
    /**
     * 寻找主类，循环 JarFile 所有条目，加载满足条件的主类
     */
    @NotNull
    protected String findMainMethod(@NotNull JarFile source) throws Exception {
        
        try (CacheURLClassLoader classLoader = new CacheURLClassLoader(getClass().getClassLoader())) {
            // 类路径
            project.getCompileClasspathElements().stream().filter(c -> !c.endsWith(".jar")).forEach(classLoader::addClasspath);
            // 依赖路径
            getLibs().forEach(artifactLib -> classLoader.addClasspath(artifactLib.getFile().getAbsolutePath()));
            
            // 主类标识注解
            Class<? extends Annotation> mainClassAnn = (Class<? extends Annotation>) classLoader.loadClass(MAIN_CLASS_ANNOTATION);
            
            // kotlin 元数据注解
            Class<? extends Annotation> kotlinMetadata = null;
            try {
                kotlinMetadata = (Class<? extends Annotation>) classLoader.loadClass("kotlin.Metadata");
            } catch (ClassNotFoundException ignore) {
            }
            
            List<String> mainClasses = new ArrayList<>();
            
            // 循环 JarFile 所有条目
            Enumeration<JarEntry> entries = source.entries();
            while (entries.hasMoreElements()) {
                JarEntry jarEntry = entries.nextElement();
                String entryName = jarEntry.getName();
                if (!entryName.endsWith(".class")) {
                    continue;
                }
                
                entryName = entryName.replace(".class", "");
                entryName = entryName.replace("/", ".");
                
                try {
                    Class<?> clazz = classLoader.loadClass(entryName);
                    if (AnnotatedElementUtils.hasAnnotation(clazz, mainClassAnn)) {
                        
                        // 如果是 kotlin 类加载 Kt 类
                        if (kotlinMetadata != null && clazz.getDeclaredAnnotation(kotlinMetadata) != null) {
                            try {
                                clazz = classLoader.loadClass(entryName + "Kt");
                            } catch (Exception ignore) {
                            }
                        }
                        
                        Method method = ReflectUtil.getMethod(clazz, "main", String[].class);
                        if (Objects.nonNull(method)) {
                            mainClasses.add(clazz.getCanonicalName());
                        }
                    }
                } catch (Exception ignore) {
                }
            }
            
            if (mainClasses.size() != 1) {
                throw new IllegalStateException(String.format("找到零个或者多个满足条件的主类：%s", mainClasses));
            }
            
            return mainClasses.get(0);
        }
    }
    
    /**
     * 获取输出的 jar 文件
     */
    protected File getTargetFile(@NotNull String finalName, @NotNull File outputDirectory) {
        if (!outputDirectory.exists()) {
            outputDirectory.mkdirs();
        }
        return new File(outputDirectory, finalName + "." + this.project.getArtifact().getArtifactHandler().getExtension())
                .getAbsoluteFile();
    }
    
    /**
     * 解析输出时间戳
     */
    @Nullable
    private FileTime parseOutputTimestamp() {
        if (this.outputTimestamp == null || this.outputTimestamp.length() < 2) {
            return null;
        }
        
        long timestampEpochSeconds;
        try {
            timestampEpochSeconds = Long.parseLong(this.outputTimestamp);
        } catch (NumberFormatException ex) {
            timestampEpochSeconds = OffsetDateTime.parse(this.outputTimestamp).toInstant().getEpochSecond();
        }
        return FileTime.from(timestampEpochSeconds, TimeUnit.SECONDS);
    }
}
